import core
import math
import copy

model_wrapper=core.model_wrapper
pyfitter=core.pyfitter
pychisq=core.pychisq
pyvec=core.pyvec

model_holder=[]

class pymodel(model_wrapper):
    def do_get_type_name(self):
        return type(self).__name__


    def do_clone(self):
        model_holder.append(self)
        return self

    def do_destroy(self):
        #print self
        for i in range(0,len(model_holder)):
            if model_holder[i] is self:
         #       print 'Loaded model:',model_holder[i],'deleted'
                del model_holder[i]
                break

    def __add__(self,r):
        return add_model(self,r)

    def __mul__(self,r):
        return mul_model(self,r)
pass        


class add_model(pymodel):
    def __init__(self,m1,m2):
        pymodel.__init__(self)
        np1=m1.get_num_params();
        np2=m2.get_num_params();
        for i in range(0,np1):
            pymodel.push_param_info(self,m1.get_param_info(i).get_name()+"1",m1.get_param_info(i).get_value())
        for i in range(0,np2):
            pymodel.push_param_info(self,m2.get_param_info(i).get_name()+"2",m2.get_param_info(i).get_value())

        self.m1=m1
        self.m2=m2

    def do_eval(self,x,p):
        p1=pyvec(self.m1.get_num_params())
        p2=pyvec(self.m2.get_num_params())

        np1=self.m1.get_num_params()
        for i in range(0,self.m1.get_num_params()):
            p1[i]=p[i]
        for i in range(0,self.m2.get_num_params()):
            p2[i]=p[i+np1]

        return self.m1.eval(x,p1)+self.m2.eval(x,p2)
pass

class mul_model(pymodel):
    def __init__(self,m1,m2):
        pymodel.__init__(self)
        np1=m1.get_num_params();
        np2=m2.get_num_params();
        for i in range(0,np1):
            pymodel.push_param_info(self,m1.get_param_info(i).get_name()+"1",m1.get_param_info(i).get_value())
        for i in range(0,np2):
            pymodel.push_param_info(self,m2.get_param_info(i).get_name()+"2",m2.get_param_info(i).get_value())

        self.m1=m1
        self.m2=m2

    def do_eval(self,x,p):
        p1=pyvec(self.m1.get_num_params())
        p2=pyvec(self.m2.get_num_params())

        np1=self.m1.get_num_params()
        for i in range(0,self.m1.get_num_params()):
            p1[i]=p[i]
        for i in range(0,self.m2.get_num_params()):
            p2[i]=p[i+np1]

        return self.m1.eval(x,p1)*self.m2.eval(x,p2)
pass



class func_wrapper_model(pymodel):
    def __init__(self,func):
        pymodel.__init__(self)
        np=len(func.func_code.co_varnames)-1
        self.np=np
        assert(np==len(func.func_defaults))
        for i in range(1,np+1):
            pymodel.push_param_info(self,func.func_code.co_varnames[i],func.func_defaults[i-1])
        self.func=func
    

    def do_eval(self,x,p):
        return self.func(x,*p)

    def reset(self,func):
        self.clear_param_info()
        np=len(func.func_code.co_varnames)-1
        self.np=np
        assert(np==len(func.func_defaults))
        for i in range(1,np+1):
            pymodel.push_param_info(self,func.func_code.co_varnames[i],func.func_defaults[i-1])
        self.func=func

pass


class linear(pymodel):
    def __init__(self):
        pymodel.__init__(self)
        pymodel.push_param_info(self,"k",1)
        pymodel.push_param_info(self,"b",1)

        

    def do_eval(self,x,p):
        return x*p[0]+p[1]

pass


class sbeta(pymodel):
    def __init__(self):
        pymodel.__init__(self)
        pymodel.push_param_info(self,"s0",1)
        pymodel.push_param_info(self,"rc",1)
        pymodel.push_param_info(self,"beta",.66)
        
    

    def do_eval(self,x,p):
        S0=p[0]
        r_c=p[1]
        beta=p[2]
        
        return S0*pow(1+(x*x)/(r_c*r_c),-3*beta+.5);

pass

class nbeta(pymodel):
    def __init__(self):
        pymodel.__init__(self)
        pymodel.push_param_info(self,"s0",1)
        pymodel.push_param_info(self,"rc",1)
        pymodel.push_param_info(self,"beta",.66)
        
    
    def do_eval(self,x,p):
        S0=p[0]
        r_c=p[1]
        beta=p[2]
        
        return S0*pow(1+(x*x)/(r_c*r_c),-3./2.*beta);

pass
    
class power_law(pymodel):
    def __init__(self):
        pymodel.__init__(self)
        pymodel.push_param_info(self,"Ampl",1)
        pymodel.push_param_info(self,"gamma",1)

    def do_eval(self,x,p):
        A=p[0]
        gamma=p[1]
#        print "%s %s"%(x,gamma)
        return A*x**gamma
pass



class bpl(pymodel):
    def __init__(self):
        pymodel.__init__(self)
        pymodel.push_param_info(self,"bpx",1)
        pymodel.push_param_info(self,"bpy",1)
        pymodel.push_param_info(self,"gamma1",1)
        pymodel.push_param_info(self,"gamma2",1)
        
      

    def do_eval(self,x,p):
        x_b=math.fabs(p[0])
        f_b=p[1]
        gamma1=p[2]
        gamma2=p[3]
        if x<x_b:
            return f_b*(x/x_b)**gamma1
        else:
            return f_b*(x/x_b)**gamma2

pass

class bremss(pymodel):
    def __init__(self):
        pymodel.__init__(self)
        pymodel.push_param_info(self,"norm",1)
        pymodel.push_param_info(self,"kT",1)
        
      

    def do_eval(self,x,p):
        norm=p[0]
        kT=p[1]

        return norm*kT**.5*math.exp(-x/kT)
pass


class gauss(pymodel):
    """
    Gaussian model
    y=N*exp(-(x-x0)^2/sigma^2)
    """
    def __init__(self):
        pymodel.__init__(self)
        pymodel.push_param_info(self,"N",1)
        pymodel.push_param_info(self,"x0",1)
        pymodel.push_param_info(self,"sigma",1)

  

    def do_eval(self,x,p):
        N=p[0]
        x0=p[1]
        sigma=p[2]
        y=(x-x0)/sigma
        return N*math.exp(-y*y/2)

pass


class norm_gauss(pymodel):
    """
    Normal distribution
    y=1/sigma/sqrt(2*pi)*exp(-(x-x0)^2/2sigma^2)
    """
    def __init__(self):
        pymodel.__init__(self)
        pymodel.push_param_info(self,"x0",1)
        pymodel.push_param_info(self,"sigma",1)

    def do_eval(self,x,p):
        x0=p[0]
        sigma=p[1]
        y=(x-x0)/sigma
        return 1/sigma/math.sqrt(math.pi)*math.exp(-y*y/2)

pass


class norm_2gauss(pymodel):
    """
    2-gaussian distribution
    """
    def __init__(self):
        pymodel.__init__(self)
        pymodel.push_param_info(self,"angle",0);
        pymodel.push_param_info(self,"x01",1)
        pymodel.push_param_info(self,"sigma1",1)
        pymodel.push_param_info(self,"x02",1);
        pymodel.push_param_info(self,"sigma2",1);
        

    def do_eval(self,x,p):
        theta=p[0]
        x01=p[1]
        sigma1=p[2]
        x02=p[3]
        sigma2=p[4]
        y1=(x-x01)/sigma1
        y2=(x-x02)/sigma2
        k1=math.cos(theta)**2
        k2=math.sin(theta)**2
        return k1/sigma1/math.sqrt(math.pi)*math.exp(-y1*y1/2)+k2/sigma2/math.sqrt(math.pi)*math.exp(-y2*y2/2)

pass


class nfw(pymodel):
    def __init__(self):
        pymodel.__init__(self)
        pymodel.push_param_info(self,"rho0",1.)
        pymodel.push_param_info(self,"rs",1.)

      

    def do_eval(self,x,p):
        rho0=p[0]
        rs=p[1]

        return rho0/(x/rs*(1+x/rs)**2)

pass


class xn(pymodel):
    def __init__(self,n):
        pymodel.__init__(self)
        self.n=n
        pymodel.push_param_info(self,"a",1)

    def do_eval(self,x,p):
        a=p[0]
        return x**self.n*a
pass


class constant(pymodel):
    def __init__(self):
        pymodel.__init__(self)
        pymodel.push_param_info(self,"c",0)
        
    def do_eval(self,x,p):
        return p[0]
pass

class front(pymodel):
    def __init__(self):
        pymodel.__init__(self)
        pymodel.push_param_info(self,"ampl1",1)
        pymodel.push_param_info(self,"gamma1",1)
        pymodel.push_param_info(self,"ampl2",1)
        pymodel.push_param_info(self,"gamma2",1)
        pymodel.push_param_info(self,"x0",1)


    def do_eval(self,x,p):
        if x<p[4]:
            return p[0]*x**p[1]
        else:
            return p[2]*x**p[3]
pass



#model_map={}



